# -*- coding: utf-8 -*-
"""
A股基本面动量策略
结合估值指标与价格动量筛选优质标的，包含A股特有机制处理
"""

from zipline.api import order_target_percent, symbol, get_datetime
from strategies.utils.a_share_utils import get_csi300_constituents, is_price_limit, apply_t1_restriction
from zipline.finance import commission, slippage
import pandas as pd
import numpy as np

def initialize(context):
    """
    初始化策略
    """
    # 设置策略参数
    context.valuation_weight = 0.4  # 估值因子权重
    context.quality_weight = 0.3    # 质量因子权重
    context.momentum_weight = 0.3   # 动量因子权重
    context.max_sector_deviation = 0.03  # 行业偏离度上限
    context.max_position = 0.05     # 单票最大仓位
    
    # 设置再平衡频率 (每月)
    context.rebalance_frequency = 21  # 约1个月
    context.days_until_rebalance = 0
    
    # 设置交易费用 (A股标准)
    context.set_commission(commission.PerShare(cost=0.0003, min_trade_cost=5))
    context.set_slippage(slippage.FixedSlippage(spread=0.001))

def handle_data(context, data):
    """
    每日交易逻辑
    """
    # 检查再平衡日
    context.days_until_rebalance -= 1
    if context.days_until_rebalance <= 0:
        rebalance_portfolio(context, data)
        context.days_until_rebalance = context.rebalance_frequency

def rebalance_portfolio(context, data):
    """
    月度再平衡逻辑
    """
    # 获取沪深300成分股
    stocks = get_csi300_constituents(context, data)
    
    # 计算因子得分
    scores = {}
    for stock in stocks:
        # 跳过涨跌停股票
        if is_price_limit(stock, data):
            continue
            
        # 计算估值因子 (PE, PB, 股息率)
        pe = data.current(stock, 'pe_ratio')
        pb = data.current(stock, 'pb_ratio')
        dy = data.current(stock, 'dividend_yield')
        valuation_score = (1/pe + 1/pb + dy) / 3
        
        # 计算质量因子 (ROE, 毛利率, 负债率)
        roe = data.current(stock, 'roe')
        gross_margin = data.current(stock, 'gross_margin')
        debt_ratio = 1 - data.current(stock, 'debt_to_assets')  # 负债率越低越好
        quality_score = (roe + gross_margin + debt_ratio) / 3
        
        # 计算动量因子 (60日收益率)
        momentum = data.history(stock, 'close', 60, '1d').pct_change()[-1]
        
        # 综合得分
        total_score = (
            context.valuation_weight * valuation_score +
            context.quality_weight * quality_score +
            context.momentum_weight * momentum
        )
        scores[stock] = total_score
    
    # 按得分排序
    sorted_stocks = sorted(scores.keys(), key=lambda k: scores[k], reverse=True)
    
    # 行业配置限制
    sector_allocation = {}
    for stock in sorted_stocks[:30]:  # 取前30名
        sector = data.current(stock, 'sector')
        sector_allocation[sector] = sector_allocation.get(sector, 0) + 1
    
    # 执行交易
    for stock in sorted_stocks[:30]:
        # 检查行业偏离
        sector = data.current(stock, 'sector')
        sector_weight = sector_allocation[sector] / 30
        if sector_weight > context.max_sector_deviation:
            continue  # 跳过超过行业偏离度的股票
            
        # 应用T+1限制
        closeable_amount = apply_t1_restriction(context, stock)
        if closeable_amount == 0:  # 无持仓时可买入
            order_target_percent(stock, context.max_position)

def analyze(context, perf):
    """
    回测结果分析
    """
    # 输出策略表现
    print("策略年化收益率: %.2f%%" % (perf.returns.mean() * 252 * 100))
    print("最大回撤: %.2f%%" % (perf.max_drawdown() * 100))
    print("夏普比率: %.2f" % perf.sharpe_ratio)
    
    # 绘制净值曲线
    perf.portfolio_value.plot()